package com.li.dao;

import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.ChatMessageDeserializer;
import dev.langchain4j.data.message.ChatMessageSerializer;
import dev.langchain4j.store.memory.chat.ChatMemoryStore;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.ValueOperations;
import org.springframework.data.redis.core.HashOperations;

import java.util.*;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;

import static dev.langchain4j.data.message.ChatMessageDeserializer.messagesFromJson;
import static dev.langchain4j.data.message.ChatMessageSerializer.messagesToJson;

/**
 * KEYS:
 *
 * memoryId:abc123 => Hash
 *     msg:001 => "{\"type\":\"HumanMessage\",\"text\":\"你好\"}"
 *     msg:002 => "{\"type\":\"AiMessage\",\"text\":\"你好！我是Qwen。\"}"
 *
 * memoryId:abc123:counter => String
 *     "2"
 */
public class RedisChatMemoryStore implements ChatMemoryStore {

    private final RedisTemplate<String, String> redisTemplate; // 推荐统一使用 String 类型
    private final String redisKeyPrefix;
    private final long expirationTimeInSeconds;

    public RedisChatMemoryStore(RedisTemplate<String, String> redisTemplate,
                                String keyPrefix,
                                long expirationTimeInSeconds) {
        this.redisTemplate = redisTemplate;
        this.redisKeyPrefix = keyPrefix;
        this.expirationTimeInSeconds = expirationTimeInSeconds;
    }

    private String keyFor(Object memoryId) {
        return redisKeyPrefix + ":" + memoryId;
    }

    @Override
    public List<ChatMessage> getMessages(Object memoryId) {
        String key = keyFor(memoryId);
        return  ChatMessageDeserializer
                .messagesFromJson(redisTemplate.opsForValue().get(key));
    }

    @Override
    public void updateMessages(Object memoryId, List<ChatMessage> messages) {
        String key = keyFor(memoryId);
        redisTemplate.opsForValue().set(key,
                ChatMessageSerializer.messagesToJson(messages),
                expirationTimeInSeconds,
                TimeUnit.SECONDS);
    }


    @Override
    public void deleteMessages(Object memoryId) {
        redisTemplate.delete(keyFor(memoryId));
    }
}
